#![allow(unused_variables)]

use sm4::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyInit};

type Sm4EcbEnc = ecb::Encryptor<sm4::Sm4>;
type Sm4EcbDec = ecb::Decryptor<sm4::Sm4>;

use hex::{FromHex};

use udf::prelude::*;

/// No data to persist
#[derive(Debug)]
struct Sm4Decrypt;

#[derive(Debug)]
struct Sm4Encrypt;


#[register]
impl BasicUdf for Sm4Encrypt {
    type Returns<'a> = Option<String>
    where
        Self: 'a;

    fn init(cfg: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
        if args.len() != 2 {
            return Err(format!("Expected 2 argument; got {}", args.len()));
        }

        let arg_key = args.get(0).unwrap().value();
        let arg_plaintext = args.get(1).unwrap().value();

        if !arg_key.is_string() {
            return Err(format!(
                "Expected string argument; got {}",
                arg_key.display_name()
            ));
        }

        let Some(arg_key_str) = arg_key.as_string() else {
            return Err(format!(
                "Expected string argument"
            ));
        };

        if arg_key_str.len() != 16 {
            return Err("The key must be a 16-byte character string".to_string());
        }

        if !arg_plaintext.is_string() {
            return Err(format!(
                "Expected string argument; got {}",
                arg_plaintext.display_name()
            ));
        }

        // max ipv6 address with colons
        // cfg.set_max_len(IPV6_MAX_LEN);

        cfg.set_maybe_null(true);

        Ok(Self)
    }

    fn process<'a>(
        &'a mut self,
        cfg: &UdfCfg<Process>,
        args: &ArgList<Process>,
        error: Option<NonZeroU8>,
    ) -> Result<Self::Returns<'a>, ProcessError> {
        let arg_key = args.get(0).unwrap().value();

        let Some(key) = arg_key.as_string() else {
            return Err(ProcessError);
        };

        if key.len() != 16 {
            return Err(ProcessError);
        }

        let arg_plaintext = args.get(1).unwrap().value();

        let Some(plaintext) = arg_plaintext.as_string() else {
            return Err(ProcessError);
        };

        let Some(r) = my_sm4_encrypt(key, plaintext) else {
            return Ok(None);
        };

        Ok(Some(r))
    }
}




#[register]
impl BasicUdf for Sm4Decrypt {
    type Returns<'a> = Option<String>
    where
        Self: 'a;

    fn init(cfg: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
        if args.len() != 2 {
            return Err(format!("Expected 2 argument; got {}", args.len()));
        }

        let arg_key = args.get(0).unwrap().value();
        let arg_ciphertext = args.get(1).unwrap().value();


        if !arg_key.is_string() {
            return Err(format!(
                "Expected string argument; got {}",
                arg_key.display_name()
            ));
        }

        let Some(arg_key_str) = arg_key.as_string() else {
            return Err(format!(
                "Expected string argument"
            ));
        };

        if arg_key_str.len() != 16 {
            return Err("The key must be a 16-byte character string".to_string());
        }

        if !arg_ciphertext.is_string() {
            return Err(format!(
                "Expected string argument; got {}",
                arg_ciphertext.display_name()
            ));
        }

        // max ipv6 address with colons
        // cfg.set_max_len(IPV6_MAX_LEN);

        cfg.set_maybe_null(true);

        Ok(Self)
    }

    fn process<'a>(
        &'a mut self,
        cfg: &UdfCfg<Process>,
        args: &ArgList<Process>,
        error: Option<NonZeroU8>,
    ) -> Result<Self::Returns<'a>, ProcessError> {
        let arg_key = args.get(0).unwrap().value();

        let Some(key) = arg_key.as_string() else {
            return Err(ProcessError);
        };

        if key.len() != 16 {
            return Err(ProcessError);
        }
        
        let arg_ciphertext = args.get(1).unwrap().value();

        let Some(ciphertext) = arg_ciphertext.as_string() else {
            return Err(ProcessError);
        };

        let Some(r) = my_sm4_decrypt(key, ciphertext) else {
            // return Err(ProcessError);
            return Ok(None);
        };

        Ok(Some(r))
    }
}


fn my_sm4_encrypt(key: &str, plaintext: &str) -> Option<String> {
    let key_bytes = key.as_bytes();
    let plaintext_bytes = plaintext.as_bytes(); 
    let res = Sm4EcbEnc::new(key_bytes.into())
        .encrypt_padded_vec_mut::<Pkcs7>(&plaintext_bytes);

    let mut hex_string = String::new();
    for byte in res {
        hex_string.push_str(&format!("{:02x}", byte));
    }
    Some(hex_string)
}

fn my_sm4_decrypt(key: &str, ciphertext: &str) -> Option<String> {
    let key_bytes = key.as_bytes();

    let conversion_result = Vec::from_hex(ciphertext);
    if let Err(_err) = conversion_result {
        return None
    }
    let ciphertext_bytes = conversion_result.unwrap();

    let res = Sm4EcbDec::new(key_bytes.into())
        .decrypt_padded_vec_mut::<Pkcs7>(&ciphertext_bytes)
        .unwrap();
    let hexstr = String::from_utf8_lossy(&res);
    Some(hexstr.to_string())
}


#[cfg(test)]
mod tests {
    use udf::mock::*;

    use super::*;


    #[test]
    fn test_init_wrong_arg_count() {
        let mut mock_cfg = MockUdfCfg::new();
        // let mut mock_args = mock_args![("localhost", "attr1", true), ("localhost", "attr2", true)];
        let mut mock_args = mock_args![("1234567890123456", "attr1", false),("b083dcc0a9f64bd9fae2fa8c936e3d776c88f739af2a29a735381f5677badef7", "attr2", false),];

        let res = Sm4Decrypt::init(mock_cfg.as_init(), mock_args.as_init());

        // assert_eq!(res.unwrap_err(), "Expected 1 argument; got 2");
        dbg!(res);
    }

    #[test]
    fn process_sm4_encrypt() {
        // Test with some random arguments
        let mut inited = Sm4Encrypt;
        let mut mock_cfg = MockUdfCfg::new();
        let mut mock_args = mock_args![("1234567890123456", "attr1", false), ("hello1", "attr2", false),];

        let res = Sm4Encrypt::process(
            &mut inited,
            mock_cfg.as_process(),
            mock_args.as_process(),
            None,
        );

        dbg!(res);

        // Our result can be weird, so we just check it has a colon
        assert!(true);
    }

    
    #[test]
    fn process_sm4_decrypt() {
        // Test with some random arguments
        let mut inited = Sm4Decrypt;
        let mut mock_cfg = MockUdfCfg::new();
        let mut mock_args = mock_args![("1234567890123456", "attr1", false),("e8ed75771fe34277a3808aaa16511936", "attr2", false),];

        let res = Sm4Decrypt::process(
            &mut inited,
            mock_cfg.as_process(),
            mock_args.as_process(),
            None,
        );

        dbg!(res);

        // Our result can be weird, so we just check it has a colon
        assert!(true);
    }

}